#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/5/8 8:55
# @Author  : LiShan
# @Email   : lishan_1997@126.com
# @File    : test.py
# @Note    : this is note
import os
import threading
import time
import psutil
import pandas as pd
import matplotlib.pyplot as plt
import numpy as  np
from winreg import OpenKey, QueryValueEx, HKEY_CURRENT_USER
import agent.dqn as ag
import torch
from environment.vissim import VisCom

device = 'cuda' if torch.cuda.is_available() else 'cpu'
ZEROREWARD = 50
project_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
net_path =  project_path.replace("\\", "/") + '/resource/vissim/3.inp'
simulation = [999999, 0, 1, 1, 42, False, False]
timming = [26, 32, 2, 130, 20, 3, [2, 2, 2, 2]]
online_net = './online_network_best.pkl'
target_net = './target_network_best.pkl'
result_file = './test_record.txt'
plans_file = './test_plans.txt'
program_name = 'run vissim.bat'


def get_desktop_path():
    key = OpenKey(HKEY_CURRENT_USER, r'Software\Microsoft\Windows\CurrentVersion\Explorer\Shell Folders')
    return QueryValueEx(key, "Desktop")[0]


def detect_vissim():
    pids = psutil.pids()
    for pid in pids:
        p = psutil.Process(pid)
        process_name = p.name()
        if "vissim" in process_name:
            return True
    else:
        return False


def normalization_process(observation):
    """flow + speed + queue"""
    # state = observation
    flow = list(map(lambda x: round((x - 0) / (25 - 0), 2), observation[0:8]))
    speed = list(map(lambda x: round((x - 0) / (60 - 0), 2), observation[8:16]))
    queue = list(map(lambda x: round((x - 0) / (400 - 0), 2), observation[16:24]))
    state = flow + speed + queue
    return state


def get_reward(delay):
    if delay < ZEROREWARD - 10:
        return 10
    elif delay < ZEROREWARD - 9:
        return 9
    elif delay < ZEROREWARD - 8:
        return 8
    elif delay < ZEROREWARD - 7:
        return 7
    elif delay < ZEROREWARD - 6:
        return 6
    elif delay < ZEROREWARD - 5:
        return 5
    elif delay < ZEROREWARD - 4:
        return 4
    elif delay < ZEROREWARD - 3:
        return 3
    elif delay < ZEROREWARD - 2:
        return 2
    elif delay < ZEROREWARD - 1:
        return 1
    elif delay < ZEROREWARD:
        return 0
    elif delay < ZEROREWARD - (-1):
        return -1
    elif delay < ZEROREWARD - (-2):
        return -2
    elif delay < ZEROREWARD - (-3):
        return -3
    elif delay < ZEROREWARD - (-4):
        return -4
    elif delay < ZEROREWARD - (-5):
        return -5
    elif delay < ZEROREWARD - (-6):
        return -6
    elif delay < ZEROREWARD - (-7):
        return -7
    elif delay < ZEROREWARD - (-8):
        return -8
    elif delay < ZEROREWARD - (-9):
        return -9
    else:
        return -10


def create_plans(timming):
    # 配时范围
    green_low = timming[0]
    green_high = timming[1]
    green_interval = timming[2]
    cycle_time = timming[3]
    loss_time = timming[4]
    amber_time = timming[5]
    clearing_time = timming[6]
    plans = []
    try:
        with open(plans_file, "a+") as f:
            f.truncate(0)
    except (IndexError, Exception):
        pass
    for a in range(green_low, green_high + 1, green_interval):
        for b in range(green_low, green_high + 1, green_interval):
            for c in range(green_low, green_high + 1, green_interval):
                for d in range(green_low, green_high + 1, green_interval):
                    if (a + b + c + d) == (cycle_time - loss_time):
                        plan = [cycle_time, amber_time, clearing_time, [a, b, c, d]]
                        plans.append(plan)
                        with open(plans_file, "a+") as f:
                            line = "%s\n" % (str(plan))
                            f.write(line)
    return plans


def draw_result():
    names = ["plan", "step", "delay", "reward"]
    data = pd.read_csv('./test_fix_record.txt', error_bad_lines=False, sep="\s+", names=names)
    delay = list(data["delay"].values)
    fix_min_idx = delay.index(min(delay))

    names = ["step", "state", "delay", "reward"]
    data = pd.read_csv('./test_fix_record_' + str(fix_min_idx) + '.txt', error_bad_lines=False, sep="\s+", names=names)
    fix_min_delay = list(data["delay"].values)

    names = ["step", "state", "delay", "reward"]
    data = pd.read_csv(result_file, error_bad_lines=False, sep="\s+", names=names)
    dqn_delay = list(data["delay"].values)
    """延误"""
    # 绘图
    x = np.linspace(0, len(dqn_delay)-1, len(dqn_delay))
    plt.plot(x, fix_min_delay, color="g")
    plt.plot(x, dqn_delay, color="r")
    plt.plot(x, np.full(len(x), np.mean(fix_min_delay)), color="b")
    plt.plot(x, np.full(len(x), np.mean(dqn_delay)), color="k")
    # 设置坐标轴范围
    # plt.xlim([-1, 501])
    # plt.ylim([35, 80])
    # 设置坐标轴刻度
    # plt.xticks(range(0, 20, 1))
    # plt.yticks(np.arange(35, 85, 5))
    # 设置坐标轴名称
    plt.xlabel("Step", fontproperties="Times New Roman", size=10.5)
    plt.ylabel("Delay", fontproperties="Times New Roman", size=10.5)
    # 设置网格
    # plt.grid()
    # 设置图例
    legend = ["Fix dalay", "DQN dalay", "Fix mean dalay", "DQN mean dalay"]
    plt.legend(legend, loc="best", frameon=False)
    # 设置标题
    plt.title("The Delay Curve", fontproperties="Times New Roman", size=10.5)
    # 保存图片
    plt.savefig("./delay_compare.svg")
    # 关闭绘图
    plt.close()
    print("绘图完成...")


def test():
    # 清空记录文件内容
    try:
        with open(result_file, 'a') as f:
            f.truncate()
    except:
        pass

    # 加载网络
    agent = ag.Agent()
    agent.online_net.load_state_dict(torch.load(online_net, map_location=device))
    agent.target_net.load_state_dict(torch.load(target_net, map_location=device))

    # 初始化参数
    test_start = time.perf_counter()
    test_mean_reward = 0
    test_mean_delay = 0
    test_step = 500

    # 获取vissim程序文件路径
    vissim_program_path = get_desktop_path().replace("\\", "/") + "/" + program_name
    # 检测vissim是否运行
    flag = detect_vissim()
    if flag:
        os.system("taskkill /F /IM vissim.exe")
        time.sleep(3)
    # 自动化启动vissim软件
    start_vissim_application = threading.Thread(target=os.startfile, args=(vissim_program_path,))
    start_vissim_application.daemon = 1
    start_vissim_application.start()
    while True:
        flag = detect_vissim()
        if flag:
            time.sleep(3)
            break
        else:
            time.sleep(1)
            continue

    # 创建配时方案
    plans = create_plans(timming)
    # 初始化vissim环境
    while True:
        try:
            env = VisCom(net_path, vissim_program_path, simulation, plans)
            break
        except:
            print("连接失败")
            continue
    print("开始测试......")
    ag.N_STATES = env.observation_space.shape[0]
    ag.N_ACTIONS = env.action_space.n
    if isinstance(env.action_space.sample(), int):
        ag.ENV_A_SHAPE = 0
    else:
        ag.ENV_A_SHAPE = env.action_space.sample().shape

    # 可视化visism仿真图像
    # env.render()

    # 重置环境获取初始交通流状态
    observation = env.reset()
    state = normalization_process(observation)

    # 热身时间
    for i in range(1):
        observation, reward, done, info = env.step(env.action_space.sample())
        state = normalization_process(observation)

    # 仿真运行指定个周期
    for step in range(test_step):
        # 智能体由交通流状态获取配时动作方案
        action = agent.action(state, random=False)

        # vissim环境采取动作运行一周期，得到下一周期的状态信息
        observation, reward, done, info = env.step(action)
        next_state = normalization_process(observation)

        # 重定义奖励
        delay = reward
        redefine_reward = get_reward(delay)

        # 更新状态、奖励、平均延误、当前回合仿真步数
        state = next_state

        if step == 0:
            test_mean_reward = redefine_reward
            test_mean_delay = delay
        else:
            test_mean_reward = (test_mean_reward + redefine_reward) / 2
            test_mean_delay = (test_mean_delay + delay) / 2

        print("step:{}, delay:{}, reward:{}".format(step, delay, redefine_reward))

        # 保存测试记录信息
        with open(result_file, 'a+') as f:
            record = "%s\t%s\t%s\t%s\t\n" % \
                     (str(step), str(state), str(delay), str(redefine_reward))
            f.write(record)

    # 输出网络的测试奖励值和延误值
    info = 'test step: {}, test_mean_reward: {}, test_mean_delay: {}'. \
        format(test_step, round(test_mean_reward, 3), round(test_mean_delay, 3))
    print(info)

    # 输出测试时间
    test_time = time.perf_counter() - test_start
    h, ss = divmod(test_time, 3600)
    m, s = divmod(ss, 60)
    info = "complete test time: {} second, that is {} hour, {} minute, {} second".format(test_time, h, m, s)
    print(info)


if __name__ == '__main__':
    # test()
    draw_result()
